#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2025/4/9
# @USER    : Shengji He
# @File    : convutils.py
# @Software: PyCharm
# @Version  : Python-
# @TASK:
from typing import Sequence, Union
import numpy as np

__all__ = ['same_padding', 'stride_minus_kernel_padding']


def same_padding(kernel_size: Union[Sequence[int], int], dilation: Union[Sequence[int], int] = 1) -> Union[
    tuple[int, ...], int]:
    """
    Return the padding value needed to ensure a convolution using the given kernel size produces an output of the same
    shape as the input for a stride of 1, otherwise ensure a shape of the input divided by the stride rounded down.

    Raises:
        NotImplementedError: When ``np.any((kernel_size - 1) * dilation % 2 == 1)``.

    """

    kernel_size_np = np.atleast_1d(kernel_size)
    dilation_np = np.atleast_1d(dilation)

    if np.any((kernel_size_np - 1) * dilation % 2 == 1):
        raise NotImplementedError(
            f"Same padding not available for kernel_size={kernel_size_np} and dilation={dilation_np}."
        )

    padding_np = (kernel_size_np - 1) / 2 * dilation_np
    padding = tuple(int(p) for p in padding_np)

    return padding if len(padding) > 1 else padding[0]


def stride_minus_kernel_padding(kernel_size: Union[Sequence[int], int], stride: Union[Sequence[int], int]) -> Union[
    tuple[int, ...], int]:
    kernel_size_np = np.atleast_1d(kernel_size)
    stride_np = np.atleast_1d(stride)

    out_padding_np = stride_np - kernel_size_np
    out_padding = tuple(int(p) for p in out_padding_np)

    return out_padding if len(out_padding) > 1 else out_padding[0]